#!/usr/bin/env python3
"""
智能训练启动脚本
自动检测GPU数量并启动单机多GPU或单GPU训练
"""

import os
import sys
import subprocess
import torch
import argparse
from pathlib import Path

def detect_gpu_count():
    """检测可用GPU数量"""
    if not torch.cuda.is_available():
        print("❌ 未检测到CUDA，将使用CPU训练")
        return 0
    
    gpu_count = torch.cuda.device_count()
    print(f"🔍 检测到 {gpu_count} 个GPU")
    
    # 显示GPU信息
    for i in range(gpu_count):
        gpu_name = torch.cuda.get_device_name(i)
        gpu_memory = torch.cuda.get_device_properties(i).total_memory / 1024**3
        print(f"   GPU {i}: {gpu_name} ({gpu_memory:.1f}GB)")
    
    return gpu_count


def launch_training(config_file, gpu_count, extra_args=None):
    """启动训练"""
    if extra_args is None:
        extra_args = []
    
    if gpu_count == 0:
        print("🚀 启动CPU训练...")
        cmd = [sys.executable, config_file] + extra_args
    elif gpu_count == 1:
        print("🚀 启动单GPU训练...")
        cmd = [sys.executable, config_file] + extra_args
    else:
        print(f"🚀 启动多GPU训练 ({gpu_count} GPUs)...")
        cmd = [
            sys.executable, "-m", "torch.distributed.launch",
            f"--nproc_per_node={gpu_count}",
            "--use_env",
            config_file
        ] + extra_args
    
    print(f"🔧 执行命令: {' '.join(cmd)}")
    print("=" * 60)
    
    try:
        subprocess.run(cmd, check=True)
        print("=" * 60)
        print("✅ 训练完成！")
        return True
    except subprocess.CalledProcessError as e:
        print("=" * 60)
        print(f"❌ 训练失败: {e}")
        return False
    except KeyboardInterrupt:
        print("=" * 60)
        print("⏹️ 训练被用户中断")
        return False

def main():
    parser = argparse.ArgumentParser(description='智能训练启动器')
    parser.add_argument('--config', type=str, default='config_distributed.py', 
                       help='配置文件路径')
    parser.add_argument('--force_cpu', action='store_true',
                       help='强制使用CPU训练')
    parser.add_argument('--args', nargs='*', default=[],
                       help='传递给配置文件的额外参数')
    
    args = parser.parse_args()
    
    # 检查配置文件是否存在
    if not os.path.exists(args.config):
        print(f"❌ 配置文件不存在: {args.config}")
        print("💡 可用的配置文件:")
        for config_file in Path('.').glob('config*.py'):
            print(f"   - {config_file}")
        return False
    
    print("🎯 智能训练启动器")
    print("=" * 60)
    
    # 检测GPU
    if args.force_cpu:
        gpu_count = 0
        print("🔧 强制使用CPU训练")
    else:
        gpu_count = detect_gpu_count()
    
    # 显示训练配置
    print(f"📋 训练配置:")
    print(f"   配置文件: {args.config}")
    print(f"   GPU数量: {gpu_count}")
    if args.args:
        print(f"   额外参数: {' '.join(args.args)}")
    
    # 确认启动
    if gpu_count > 1:
        print(f"\n⚠️  将使用 {gpu_count} 个GPU进行分布式训练")
    
    response = input("\n🚀 是否开始训练? (y/N): ").strip().lower()
    if response not in ['y', 'yes', '是']:
        print("❌ 训练已取消")
        return False
    
    # 启动训练
    return launch_training(args.config, gpu_count, args.args)

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)
